#!/usr/bin/env python

import sys
import os
import subprocess

def cmd_get_output(cmd):
    return subprocess.check_output(cmd,shell=True)

class Parted:

    def __init__(self,devpath):
        self._devpath = devpath
        self._data = cmd_get_output('LANG=C parted ' + self._devpath+ ' print')
        self._parser()

    def _parser(self):
        vv = self._data.split('\n\n')
        self._headers_parser(vv[0])
        self._part_parser(vv[1])
        pass

    def _headers_parser(self,headers):
        vv = headers.split('\n')
        vv_dict ={}
        for it in vv:
            kv=it.strip().split(':')
            vv_dict[kv[0]]=kv[1].strip()
        self._part_type = vv_dict['Partition Table']

    def _part_parser(self,part):
        vv = part.split('\n')
        self._part_dict = []
        for it in vv:
            it = it.strip()
            lt = it.split(None,6)
            if lt[0].isdigit():
                d= {}
                d['num']  = lt[0]
                d['size'] = lt[3]
                if lt[4] == 'extended':
                    if len(lt) >= 6:
                        d['flag'] = lt[5]
                else:
                    d['fstype'] = lt[5]
                    if len(lt) >= 7:
                        d['flag'] = lt[6]
                self._part_dict.append(d)

    def _get_boot_part(self):
        boot = self._part_dict[0]
        for d in self._part_dict:
            if d.has_key('flag') and 'boot' in d['flag']:
                boot = d
                break
        return boot

    def install(self,source):
        devpath = self._devpath + self._get_boot_part()['num']
        mdir = self._is_mounted(devpath)
        to_umount = False
        if mdir is None:
            mdir = '/tmp/.tmpmountdir'
            self.mount_boot_part(devpath,mdir)
            to_umount = True
        cmd = 'cp -avr %s/* %s' % (source,mdir)
        subprocess.call(cmd,shell=True)

        isolinuxdir=os.path.join(source,'isolinux')
        if os.path.isdir(isolinuxdir):
            syslinuxdir=os.path.join(source,'syslinux')
            cmd='mkdir -p %s ;cp -av %s/isolinux.cfg %s/syslinux.cfg' %(syslinuxdir,isolinuxdir,syslinuxdir)
            subprocess.check_call(cmd,shell=True)

            cmd=['syslinux','--force','--install',devpath]
            subprocess.check_call(cmd)

        
        bootlace=os.path.join(source,'grub4dos/bootlace.com')
        if os.path.exists(bootlace):
            cmd=bootlace+' --time-out=0 --mbr-disable-floppy ' + self._devpath
            subprocess.check_call(cmd,shell=True)
            cmd='cp -av %s/grub4dos/menu.lst %s' %(source,mdir)
            subprocess.check_call(cmd,shell=True)

        if to_umount:
            cmd = 'sync;sync;sudo umount ' + mdir
            subprocess.call(cmd,shell=True)

    def _install_syslinux(self,root_dir):
        cmd=['syslinux','--install',self._devpath]
        subprocess.check_call(cmd)


    def _is_mounted(self,devpath):
        r=None
        res = None
        try:
            r = open('/proc/mounts')
            for l in r:
                vv = l.split(None,3)
                if vv[0] == devpath:
                    res = vv[1]
                    break
        except:pass
        finally:
            if r:r.close()
        return res

    def mount_boot_part(self,devpath,mdir):
        cmd='mkdir -p ' + mdir +' ; sudo umount ' + mdir
        subprocess.call(cmd,shell=True)
        cmd = 'sudo mount -o uid=%d,gid=%d %s %s' % (os.getuid(),os.getgid(),devpath,mdir)
        return subprocess.call(cmd,shell=True)


### install2disk cdroot devpath
if __name__ == '__main__':
    pt = Parted(sys.argv[2])
    pt.install(sys.argv[1])
